Repository navigation
feat: add DeepSeek V4 Flash GRPO support - Automodel path - #3595
Conversation
074a7fd to
0bcd08a
Compare
5d0a67b to
4286bb0
Compare
|
/ok to test f9511e0 |
|
/ok to test 0e43437 |
d6a3cf1 to
2020329
Compare
|
/ok to test 2020329 |
|
/ok to test 4ba8402 |
|
/ok to test 879376f |
879376f to
0a4e04a
Compare
|
/ok to test 0a4e04a |
a0748c7 to
a1d83ee
Compare
zpqiu
left a comment
There was a problem hiding this comment.
Reviewed with a team of six agents (RL/codebase, vLLM, Automodel, tests, bug-hunting, and triage of the existing threads), against vLLM 0.25.1 (752a3a50) and Automodel 1814c6c9. Every finding below was re-verified against source before posting.
Nice work — this is a large, genuinely difficult integration and it holds up well under scrutiny. Highlights from what we checked and found clean:
nemo_rl/data/deepseek_v4_encoding.pyis byte-identical to vLLM 0.25.1's copy apart from one added license line — no adaptation drift.- All three non-obvious vLLM APIs are used correctly:
replace_parameter(..., prefer_copy=True),deepgemm_post_process_fp8_weight_block(..., is_bmm=, bmm_batch_size=), and theget_layerwise_info(...).restore_metadata/SKIP_TENSORSlifecycle. The ordering aroundrecord_metadata_for_reloadingis the only window that works, and the PR hits it. - The TP1/CP8/EP128 recipe is valid and internally consistent (
ep_shard_size = 1is the correct branch at exact equality; EP128 dividesn_routed_experts: 256). - Every DeepSeek V4 test genuinely executes — 0 skipped, confirmed three ways (item-count reconciliation, the runner's nodeid manifest, and a negative control showing the marker filter bites).
ruff format,ruff check,pyrefly, andconfig_cli.py minimize-checkare all clean. - The prior round's report of 21 failing committed tests does not reproduce at this HEAD; all 21 existing review threads are stale or fixed, and we've resolved them.
Two notes on scope, so the findings below read correctly. Several are latent rather than live: the rendering issue is not on this recipe's GRPO path, and the dp_size guard is pre-existing code your recipe is simply the first to make degenerate. They're raised because they'd bite the next person, not because they affect your validated run.
One request: would you be able to share throughput numbers for the run behind the reference curves — tokens/sec and peak memory at the 16-node TP1/CP8/EP128 configuration? The convergence curves through step 50 are great for correctness, but there's no baseline for what the FP8 refit path costs per step, which is what the next person sizing a DSV4 run will want.
Generated by Claude Code
zpqiu
left a comment
There was a problem hiding this comment.
Follow-up to my earlier review with three findings that landed after I published it -- my mistake for posting before the whole team had reported back. Two of these are things I should have caught in the first pass.
data.filter_overlong_prompts: trueis read by nothing. It appears exactly once in the repo (this recipe) and is not aDataConfigfield. Same class as the deadNRL_*env vars already flagged.deepseek_v4_fp8.pyis missing frompyrefly.tomlwhile both its quantization siblings are listed, which is the only reasonpyrefly checkreports clean. Adding it surfaces one real error. This qualifies my earlier "lint is clean at HEAD" note: it is clean, but the PR's largest new module is not being typechecked.- A stride-shaped hole in the DSV4 test suite, found by mutation testing rather than inspection. The production code is correct; the suite simply would not catch one specific silent regression. A verified test is included.
Everything else from the first review stands unchanged.
Generated by Claude Code
zpqiu
left a comment
There was a problem hiding this comment.
Final follow-up -- four test-quality findings, closing out the review. Sorry for the trailing reviews; the test agent's mutation work finished after I'd already posted, and these were worth not dropping.
num_special_tokens_to_addreturns 1 instead of 0 against a real HF tokenizer. This sharpens my earlier comment on the same area: the issue is not only that the test is weak, the production method is wrong and the stub masks it.process_weights_after_loading_moehas zero test coverage repo-wide, and this PR changed how it manages parameter storage. The sibling linear path has an incident-driven guard for exactly this failure mode; the MoE path holds more memory and has none. Verified test included.- The narrowed
grpo.setup()KV-cache guard is untested on either side of the change. - A
finalize_refitnegative test cannot fail -- confirmed by mutation.
All findings verified against source; both included tests were run at HEAD and mutation-tested to confirm they are not tautologies.
To restate the overall picture, since this is the third review and the volume may read worse than the PR deserves: the implementation holds up well. The vendored encoder is byte-identical to upstream, all three non-obvious vLLM APIs are used correctly, the CP8/EP128 parallelism is valid, and every DeepSeek V4 test genuinely executes with zero skips (proven four ways). Roughly as many candidate findings were killed under scrutiny as were posted -- including two of my own.
Generated by Claude Code
|
/ok to test c7c4d76 |
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Preserve original refit errors, remove ineffective recipe settings, and cover FP8 layout, tokenizer, and KV-cache contracts without expanding model support. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
e874e79 to
0aa1b61
Compare
Reject refit paths that bypass DeepSeek V4 preparation and finalization. Require the nightly run to reach five steps with finite metrics, and pin the validated attention backend. Cover compatible FP32 scale references across two refits using real vLLM quant configs. Document colocated H100 validation, precision mismatch, Blackwell limitations, and the development branch's QAT work. Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
|
/ok to test d74d503 |
Restore copy_ semantics for other FP8 MoE models and cover model selection, parameter identity, and incompatible layouts in the refit tests. Remove the Ray V1 executor override from the DeepSeek V4 recipe and guide. Preserve upstream comments to keep the review diff focused. Validation: 134 related unit tests and applicable pre-commit hooks passed on a compute node. The RayExecutorV2 nightly completed five training steps and passed all metric checks; its trailing pre-commit wrapper failed on an unset SLURM_JOB_ID. Step 3 had a token_mult_prob_error spike of 1.795, while the median remained within the configured threshold. Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
|
/ok to test 628e0cc |
…rwise reload Two vLLM 0.29 changes silently broke the DeepSeek V4 refit path that #3595 validated on 0.25.1; both surfaced once 10e1a86 let the refit stream past its first batch. 1. vLLM 0.29 gives DeepseekV4ForCausalLM a `packed_modules_mapping` in which `wkv` is a shard of two fused modules (`fused_wqa_wkv` and `fused_wkv_wgate`). `_get_module_from_param_name` resolved a shard by its leaf name alone, so `attn.wkv` was sent to `attn.fused_wkv_wgate`, the lookup failed (`Could not find module for parameter 'model.layers.0.attn.wkv.weight'`), the weight was treated as non-fp8 and refit unquantized. Shard names that appear under more than one fused module are now left to the parent-aware DeepSeek V4 remap. 2. vLLM 0.29 split the layerwise-reload skip list: `SKIP_TENSORS` only keeps a tensor off the meta device, while the new `SKIP_LOAD_TENSORS` is what keeps its weight loader unwrapped and out of the layer's load accounting (0.25.1 did both with `SKIP_TENSORS`). The immediately loaded expert tensors were therefore counted, vLLM ran `process_weights_after_loading` on each expert layer as soon as its last expert arrived, and `finalize_refit` converted the already kernel-layout weights and scales a second time. `prepare_refit` now registers the names in both sets and returns a `SkipNames` record so `restore_refit` removes exactly what was added from each set. Also on the 0.29 test side: `WeightsMapper.get_unstacked_mapper` is `get_rename_mapper` (what vLLM's own loader passes to `apply_vllm_mapper`), `Fp8MoEMethod.__init__` reads the MoE config's partition sizes, the import-target test skips imports inside `try`/`except ImportError` version fallbacks, and the tied-alias filter tolerates module-less stand-ins (a bare `load_weights` callable has no tied embeddings to drop). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
What does this PR do?
Adds GRPO support for
deepseek-ai/DeepSeek-V4-Flash(-0731)anddeepseek-ai/DeepSeek-V4-Flash-Baseusing the AutoModel training backend and colocated vLLM generation.The PR adds:
fp8_ds_mlaKV-cache configuration and validation.Why are these changes needed?
DeepSeek V4 tokenizer
The supported checkpoints do not provide a Jinja chat template that Hugging Face
apply_chat_template()can use. DeepSeek V4 instead defines its chat formatting through a Python encoder. Because the NeMo RL driver relies onapply_chat_template(), this PR vendors vLLM 0.25.1's DeepSeek V4 encoder and adapts it to the Hugging Face tokenizer interface. With the corresponding NeMo RL and vLLM tokenizer modes enabled by the recipe, training and rollout use consistent rendering, includingenable_thinking.fp8_ds_mlafp8_ds_mlais a vLLM runtime KV-cache format rather than a value selected by the checkpoints'config.json. DeepSeek V4 sparse MLA uses a packed block-scaled layout containing FP8 NoPE state, BF16 RoPE state, and embedded block scales. This differs from the traditional per-tensor FP8 KV caches handled by NeMo RL's generick_scale/v_scalerefit path. This PR accepts the format while excluding it from per-tensor KV-scale calculation and synchronization, and validates the DeepSeek V4 configuration to keep the driver and realized vLLM cache layout consistent.FP8 refit
The validated recipe trains dequantized BF16 weights with AutoModel while the colocated vLLM rollout model uses blockwise FP8 weights. DeepSeek V4 has model-specific Hugging Face-to-vLLM name mappings, fused projections, shape-changing
wo_aBMM post-processing, and packed RoutedExperts layouts. These prevent repeated refits from directly loading checkpoint-shaped tensors into the already realized vLLM runtime parameters. The DeepSeek V4-scoped mapping and vLLM native layerwise reload lifecycle restore checkpoint-loadable layouts before each refit and rebuild the kernel layouts afterward.Issues
N/A
Usage
The validated reference configuration uses 16 nodes with 8 H100 GPUs per node:
See DeepSeek V4 Flash for supported models, configuration details, and limitations.
Validation
e874e79.Before your PR is "Ready for review"
Pre checks:
Additional Information
The validated scope is DeepSeek V4 Flash 0731 and Flash Base with AutoModel training and vLLM generation. DeepSeek V4 Pro, Megatron training, SGLang generation, and training-side FP8 fake quantization are outside the scope of this PR.